package com.mid;

import java.util.LinkedList;
import java.util.Queue;

/**
 * Created by Lxk on 2020/1/25.
 */

class TreeNextNode {

    public int val;
    public TreeNextNode left;
    public TreeNextNode right;
    public TreeNextNode next;

    public TreeNextNode() {
    }

    public TreeNextNode(int _val) {
        val = _val;
    }

    public TreeNextNode(int _val, TreeNextNode _left, TreeNextNode _right, TreeNextNode _next) {
        val = _val;
        left = _left;
        right = _right;
        next = _next;
    }
};

public class Solution117 {

    public TreeNextNode connect(TreeNextNode root) {
        if (root == null) {
            return null;
        }
        Queue<TreeNextNode> queue = new LinkedList<>();
        queue.add(root);
        TreeNextNode head = root, tail = null;
        while (!queue.isEmpty()) {
            TreeNextNode cur = queue.poll();
            if (cur.left != null) {
                queue.add(cur.left);
                tail = cur.left;
            }
            if (cur.right != null) {
                queue.add(cur.right);
                tail = cur.right;
            }
            if (cur == head) {
                int num = queue.size();
                for (int i = 1; i <= num; i++) {
                    TreeNextNode curNode = queue.poll();
                    queue.add(curNode);
                    if (i != num) {
                        curNode.next = queue.peek();
                    } else {
                        curNode.next = null;
                    }
                }
                head = tail;
            }
        }
        return root;
    }

}
